#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
HanLP安装脚本
自动安装HanLP及其依赖
"""

import os
import sys
import subprocess
import platform

def check_python_version():
    """检查Python版本"""
    if sys.version_info < (3, 6):
        print("错误: 需要Python 3.6或更高版本")
        return False
    return True

def install_package(package, upgrade=False):
    """安装Python包"""
    cmd = [sys.executable, "-m", "pip", "install"]
    if upgrade:
        cmd.append("--upgrade")
    cmd.append(package)
    
    print(f"正在安装{package}...")
    try:
        subprocess.check_call(cmd)
        return True
    except subprocess.CalledProcessError as e:
        print(f"安装{package}失败: {e}")
        return False

def setup_environment():
    """设置环境变量"""
    # 确保HANLP_HOME环境变量存在
    hanlp_home = os.environ.get("HANLP_HOME")
    if not hanlp_home:
        # 设置默认路径
        if platform.system() == "Windows":
            hanlp_home = os.path.join(os.path.expanduser("~"), ".hanlp")
        else:
            hanlp_home = os.path.join(os.path.expanduser("~"), ".hanlp")
        
        print(f"未设置HANLP_HOME环境变量，将使用默认路径: {hanlp_home}")
        
        # 创建目录（如果不存在）
        os.makedirs(hanlp_home, exist_ok=True)
        
        # 添加环境变量提示
        if platform.system() == "Windows":
            print("请添加以下环境变量:")
            print(f"setx HANLP_HOME {hanlp_home}")
        else:
            print("请添加以下环境变量到~/.bashrc或~/.zshrc:")
            print(f"export HANLP_HOME={hanlp_home}")
    
    return True

def install_requirements():
    """安装基本依赖"""
    requirements = [
        "pip>=21.0",
        "setuptools>=50.0",
        "wheel>=0.35",
        "flask>=2.0.0",
        "flask-cors>=3.0.0"
    ]
    
    for req in requirements:
        if not install_package(req, upgrade=True):
            return False
    
    return True

def install_hanlp():
    """安装HanLP及其依赖"""
    # 安装PyTorch (根据系统自动选择CPU或GPU版本)
    try:
        import torch
        print("已安装PyTorch:", torch.__version__)
    except ImportError:
        print("正在安装PyTorch...")
        if platform.system() == "Windows":
            # Windows环境通常安装CPU版本即可
            install_package("torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cpu")
        else:
            # Linux/Mac环境，使用pip自动选择
            install_package("torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2")
    
    # 安装HanLP
    if not install_package("hanlp", upgrade=True):
        return False
    
    # 验证安装
    try:
        import hanlp
        print(f"HanLP安装成功! 版本: {hanlp.__version__}")
        return True
    except ImportError as e:
        print(f"HanLP安装失败: {e}")
        return False

def create_requirements_file():
    """创建requirements.txt文件"""
    with open("requirements.txt", "w", encoding="utf-8") as f:
        f.write("# HanLP依赖\n")
        f.write("hanlp>=2.1.0\n")
        f.write("torch>=2.0.0\n")
        f.write("# Web服务依赖\n")
        f.write("flask>=2.0.0\n")
        f.write("flask-cors>=3.0.0\n")
    
    print("已创建requirements.txt文件")

def main():
    """主函数"""
    print("="*50)
    print("HanLP安装程序")
    print("="*50)
    
    # 检查Python版本
    if not check_python_version():
        return False
    
    # 安装基本依赖
    if not install_requirements():
        print("安装基本依赖失败")
        return False
    
    # 设置环境
    if not setup_environment():
        print("设置环境失败")
        return False
    
    # 安装HanLP
    if not install_hanlp():
        print("安装HanLP失败")
        return False
    
    # 创建requirements.txt
    create_requirements_file()
    
    # 安装完成
    print("\n"+"="*50)
    print("HanLP安装完成!")
    print("="*50)
    print("\n运行方式:")
    print("1. 启动API服务: python run_hanlp_server.py")
    
    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1) 